#!/usr/bin/perl -w
use strict;
use Getopt::Long;

#88888888888888888888888888888888888888888888888888888888888888888888888888888888

#	Usage

#88888888888888888888888888888888888888888888888888888888888888888888888888888888

my $usage = <<'USAGE';

USAGE:

    filter_fasta_by_regex_list.pl
                --regex "Gene:(ENS[A-Z0-9]+)"
                --list list_of_matching_IDs
				[--histogram sequence_histogram_output_file]
				[--verbose]
				[--help]
	
	Iterates through a fasta file and filters out any non-matching sequences
	IDs are extracted from the accession line using the supplied regular
	    expression and compared to the supplied list.
USAGE


#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

#	Get options

#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888
# mandatory
my $regex;
my $list;

# optional parameters
my $verbose;
my $histogram;
my $help;
GetOptions(
			'list=s'			=> \$list,
			'regex=s'			=> \$regex,
			'histogram=s'		=> \$histogram,
			'verbose'			=> \$verbose,
			'help'				=> \$help);

die $usage if ($help);
die $usage unless ($regex && $list);

#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

#	get list of matching IDs

#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

open FILE, $list or die "Error!\nCould not open the file [$list]\n$!\n";
my @file_IDs = <FILE>;
chomp @file_IDs;
close FILE;

# uniquify array of IDs and save duplicates
my %IDs;
$IDs{$_}++ foreach (@file_IDs);

my @duplicates;
my $cnt_duplicate;
foreach (keys %IDs)
{
	push(@duplicates,  $_) if $IDs{$_} > 1;
	$cnt_duplicate += ($IDs{$_}-1);
}


# set hash values to zero: this will hold the number of sequences matching that ID
$IDs{$_} = 0 foreach (keys %IDs);


die "Error!\nNo IDs found in [$list]\n" unless @file_IDs;

my $cnt_seq = 0;
my $cnt__filtered_seq = 0;
PARSE_FILE:
while (<>)
{
	
	PARSE_LINE:
	{
		if (/>/)
		{
			$cnt_seq++;
			my @accIDs = /$regex/g;
			if ($verbose && !@accIDs)
			{
				print STDERR	"Warning!\nNo IDs matching [$regex] were found on the ".
								"accession line $.\n";
			}
			foreach my $id (@accIDs)
			{
				next unless exists($IDs{$id});
				$cnt__filtered_seq++;
				++$IDs{$id};
				print;
				while (<>)
				{
					redo PARSE_LINE if (/>/);
					print;
				}
				last PARSE_FILE;
			}
		}
	}
}

print STDERR "\t$cnt_seq sequences\tin the original file(s)\n";
print STDERR "\t", scalar @file_IDs, "\tIDs in the supplied list\n";
print STDERR "\t", scalar keys %IDs, "\tunique IDs in the supplied list\n"
	if (scalar keys %IDs != scalar @file_IDs);
print STDERR "\t", scalar @duplicates, "\tIDs with duplicates in the list\n" if @duplicates;
print STDERR "\t\ti.e. ", $cnt_duplicate, "\tduplicates in the list\n" if @duplicates;

my $cnt_id_matches = 0;
for (keys %IDs)
{
	$cnt_id_matches += $IDs{$_};
}
print STDERR "\t", $cnt_id_matches, "\tIDs which matched sequences\n";
print STDERR "\t", '-'x6, "\n";
print STDERR "\t$cnt__filtered_seq\tsequences after filtering\n";

# print counts of sequences which match > 1 ID
if ($histogram)
{
	open HISTFILE, ">$histogram" or die "Error!\nCould not open the file [$histogram]\n$!\n";
	my @sorted_keys = sort {$IDs{$a} <=> $IDs{$b} } keys %IDs;
	if (@sorted_keys )
	{
		my $padding_spaces = length($sorted_keys[0]) - 2;
		print HISTFILE "ID", ' ' x $padding_spaces, "\t", "Count of matching sequences\n";
		print HISTFILE "--", ' ' x $padding_spaces, "\t", "---------------------------\n";

		foreach (@sorted_keys)
		{
			print HISTFILE "$_\t$IDs{$_}\n"
		}
	}
	
	
	if (@duplicates)
	{
		print HISTFILE "\n\n", scalar @duplicates, " duplicate IDs in the supplied list:\n";
		print HISTFILE join ("\n",  @duplicates,  "");
	}


}
